package com.gaff.emp.core.datasource.config;

import cn.hutool.core.util.StrUtil;
import com.alibaba.druid.pool.DruidDataSource;
import com.gaff.emp.core.datasource.properties.DataSourceProperties;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 多数据源
 */
public class DynamicRoutingDataSource extends AbstractRoutingDataSource {

    private DruidDataSource _defaultDataSource;

    private List<DruidDataSource> _targetDataSources;

    public void addDataSources(DataSourceProperties dataSourceProperties) {
        DruidDataSource druidDataSource = DynamicDataSourceFactory.buildDruidDataSource(dataSourceProperties);
        if (_targetDataSources == null) {
            _targetDataSources = new ArrayList<>();
        }
        _targetDataSources.add(druidDataSource);
        setTargetDataSources(_targetDataSources);
        super.afterPropertiesSet();
    }

    public void removeDataSource(String dataSourceName) {
        if (_targetDataSources == null || _targetDataSources.size() == 0) {
            return;
        }
        List<DruidDataSource> list = new ArrayList<>();
        for (DruidDataSource druidDataSource : _targetDataSources) {
            if (druidDataSource.getName().equals(dataSourceName)) {
                continue;
            }
            list.add(druidDataSource);
        }
        setTargetDataSources(list);
        super.afterPropertiesSet();

    }

    public void setTargetDataSources(List<DruidDataSource> targetDataSources) {
        if (targetDataSources == null || targetDataSources.size() == 0) {
            return;
        }
        Map<Object, Object> druidDataSourceMap = new HashMap<>();
        for (DruidDataSource druidDataSource : targetDataSources) {
            druidDataSourceMap.put(druidDataSource.getName(), druidDataSource);
        }
        super.setTargetDataSources(druidDataSourceMap);

        this._targetDataSources = targetDataSources;
        this._defaultDataSource = _targetDataSources.get(0);
    }

    @Override
    protected Object determineCurrentLookupKey() {
        String dataSourceKey = DynamicContextHolder.peek();
        if (StrUtil.isBlank(dataSourceKey)) {
            return _defaultDataSource.getName();
        }
        return dataSourceKey;
    }
}
